// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/functions/prestosql/aggregates/multi_map_aggAggregate.h>
#include <pollux/exec/addressable_non_null_value_list.h>
#include <pollux/exec/aggregate.h>
#include <pollux/exec/strings.h>
#include <pollux/functions/lib/aggregates/value_list.h>
#include <pollux/functions/prestosql/aggregates/aggregate_names.h>
#include <pollux/type/floating_point_util.h>
#include <pollux/vector/flat_vector.h>

namespace kumo::pollux::aggregate::prestosql {
    namespace {
        /// Maintains a key-to-<list-of-values> map. Keys must be non-null.
        template<
            typename K,
            typename Hash = std::hash<K>,
            typename EqualTo = std::equal_to<K> >
        struct MultiMapAccumulator {
            melon::F14FastMap<
                K,
                ValueList,
                Hash,
                EqualTo,
                AlignedStlAllocator<std::pair<const K, ValueList>, 16> >
            keys;

            Strings stringKeys;

            MultiMapAccumulator(const TypePtr & /*type*/, HashStringAllocator *allocator)
                : keys{
                    AlignedStlAllocator<std::pair<const K, ValueList>, 16>(
                        allocator)
                } {
            }

            MultiMapAccumulator(
                Hash hash,
                EqualTo equalTo,
                HashStringAllocator *allocator)
                : keys{
                    0, // initialCapacity
                    hash,
                    equalTo,
                    AlignedStlAllocator<std::pair<const K, ValueList>, 16>(
                        allocator)
                } {
            }

            size_t size() const {
                return keys.size();
            }

            size_t numValues() const {
                size_t cnt = 0;
                for (const auto &entry: keys) {
                    cnt += entry.second.size();
                }
                return cnt;
            }

            /// Adds key-value pair.
            void insert(
                const DecodedVector &decodedKeys,
                const DecodedVector &decodedValues,
                vector_size_t index,
                HashStringAllocator &allocator) {
                auto &values = insertKey(decodedKeys, index, allocator);
                values.appendValue(decodedValues, index, &allocator);
            }

            /// Adds a key with a list of values.
            void insertMultiple(
                const DecodedVector &decodedKeys,
                vector_size_t keyIndex,
                const DecodedVector &decodedValues,
                vector_size_t valueIndex,
                vector_size_t numValues,
                HashStringAllocator &allocator) {
                auto &values = insertKey(decodedKeys, keyIndex, allocator);
                for (auto i = 0; i < numValues; ++i) {
                    values.appendValue(decodedValues, valueIndex + i, &allocator);
                }
            }

            ValueList &insertKey(
                const DecodedVector &decodedKeys,
                vector_size_t index,
                HashStringAllocator &allocator) {
                auto key = decodedKeys.value_at<K>(index);
                if constexpr (std::is_same_v<K, StringView>) {
                    if (!key.isInline() && !keys.contains(key)) {
                        key = stringKeys.append(key, allocator);
                    }
                }

                return keys.insert({key, ValueList()}).first->second;
            }

            void extract(
                VectorPtr &mapKeys,
                ArrayVector &mapValueArrays,
                vector_size_t &keyOffset,
                vector_size_t &valueOffset) {
                auto flatKeys = mapKeys->as_flat_vector<K>();
                auto mapValues = mapValueArrays.elements();

                for (auto &entry: keys) {
                    flatKeys->set(keyOffset, entry.first);

                    const auto numValues = entry.second.size();
                    mapValueArrays.setOffsetAndSize(keyOffset, valueOffset, numValues);

                    aggregate::ValueListReader reader(entry.second);
                    for (auto i = 0; i < numValues; i++) {
                        reader.next(*mapValues, valueOffset++);
                    }

                    ++keyOffset;
                }
            }

            void free(HashStringAllocator &allocator) {
                for (auto &entry: keys) {
                    entry.second.free(&allocator);
                }

                stringKeys.free(allocator);
            }
        };

        struct ComplexTypeMultiMapAccumulator {
            MultiMapAccumulator<
                AddressableNonNullValueList::Entry,
                AddressableNonNullValueList::Hash,
                AddressableNonNullValueList::EqualTo>
            base;

            /// Stores unique non-null keys.
            AddressableNonNullValueList serializedKeys;

            ComplexTypeMultiMapAccumulator(
                const TypePtr &type,
                HashStringAllocator *allocator)
                : base{
                    AddressableNonNullValueList::Hash{},
                    AddressableNonNullValueList::EqualTo{type},
                    allocator
                } {
            }

            size_t size() const {
                return base.size();
            }

            size_t numValues() const {
                return base.numValues();
            }

            /// Adds key-value pair.
            void insert(
                const DecodedVector &decodedKeys,
                const DecodedVector &decodedValues,
                vector_size_t index,
                HashStringAllocator &allocator) {
                const auto entry = serializedKeys.append(decodedKeys, index, &allocator);

                auto &values = insertKey(entry);
                values.appendValue(decodedValues, index, &allocator);
            }

            /// Adds a key with a list of values.
            void insertMultiple(
                const DecodedVector &decodedKeys,
                vector_size_t keyIndex,
                const DecodedVector &decodedValues,
                vector_size_t valueIndex,
                vector_size_t numValues,
                HashStringAllocator &allocator) {
                const auto entry = serializedKeys.append(decodedKeys, keyIndex, &allocator);

                auto &values = insertKey(entry);
                for (auto i = 0; i < numValues; ++i) {
                    values.appendValue(decodedValues, valueIndex + i, &allocator);
                }
            }

            ValueList &insertKey(const AddressableNonNullValueList::Entry &key) {
                auto result = base.keys.insert({key, ValueList()});
                if (!result.second) {
                    serializedKeys.removeLast(key);
                }

                return result.first->second;
            }

            void extract(
                VectorPtr &mapKeys,
                ArrayVector &mapValueArrays,
                vector_size_t &keyOffset,
                vector_size_t &valueOffset) {
                auto &mapValues = mapValueArrays.elements();

                for (auto &entry: base.keys) {
                    AddressableNonNullValueList::read(entry.first, *mapKeys, keyOffset);

                    const auto numValues = entry.second.size();
                    mapValueArrays.setOffsetAndSize(keyOffset, valueOffset, numValues);

                    aggregate::ValueListReader reader(entry.second);
                    for (auto i = 0; i < numValues; i++) {
                        reader.next(*mapValues, valueOffset++);
                    }

                    ++keyOffset;
                }
            }

            void free(HashStringAllocator &allocator) {
                base.free(allocator);
                serializedKeys.free(allocator);
            }
        };

        // A wrapper around MultiMapAccumulator that overrides hash and equal_to
        // functions to use the custom comparisons provided by a custom type.
        template<TypeKind Kind>
        struct CustomComparisonMultiMapAccumulator {
            using NativeType = typename TypeTraits<Kind>::NativeType;

            struct Hash {
                const TypePtr &type;

                size_t operator()(const NativeType &value) const {
                    return static_cast<const CanProvideCustomComparisonType<Kind> *>(
                                type.get())
                            ->hash(value);
                }
            };

            struct EqualTo {
                const TypePtr &type;

                bool operator()(const NativeType &left, const NativeType &right) const {
                    return static_cast<const CanProvideCustomComparisonType<Kind> *>(
                               type.get())
                           ->compare(left, right) == 0;
                }
            };

            // The underlying MultiMapAccumulator to which all operations are
            // delegated.
            MultiMapAccumulator<
                NativeType,
                CustomComparisonMultiMapAccumulator::Hash,
                CustomComparisonMultiMapAccumulator::EqualTo>
            base;

            CustomComparisonMultiMapAccumulator(
                const TypePtr &type,
                HashStringAllocator *allocator)
                : base{
                    CustomComparisonMultiMapAccumulator::Hash{type},
                    CustomComparisonMultiMapAccumulator::EqualTo{type},
                    allocator
                } {
            }

            size_t size() const {
                return base.size();
            }

            size_t numValues() const {
                return base.numValues();
            }

            // Adds key-value pair.
            void insert(
                const DecodedVector &decodedKeys,
                const DecodedVector &decodedValues,
                vector_size_t index,
                HashStringAllocator &allocator) {
                base.insert(decodedKeys, decodedValues, index, allocator);
            }

            // Adds a key with a list of values.
            void insertMultiple(
                const DecodedVector &decodedKeys,
                vector_size_t keyIndex,
                const DecodedVector &decodedValues,
                vector_size_t valueIndex,
                vector_size_t numValues,
                HashStringAllocator &allocator) {
                base.insertMultiple(
                    decodedKeys, keyIndex, decodedValues, valueIndex, numValues, allocator);
            }

            ValueList &insertKey(
                const DecodedVector &decodedKeys,
                vector_size_t index,
                HashStringAllocator &allocator) {
                return base.insertKey(decodedKeys, index, allocator);
            }

            void extract(
                VectorPtr &mapKeys,
                ArrayVector &mapValueArrays,
                vector_size_t &keyOffset,
                vector_size_t &valueOffset) {
                base.extract(mapKeys, mapValueArrays, keyOffset, valueOffset);
            }

            void free(HashStringAllocator &allocator) {
                base.free(allocator);
            }
        };

        template<typename T>
        struct MultiMapAccumulatorTypeTraits {
            using AccumulatorType = MultiMapAccumulator<T>;
        };

        // Ensure Accumulator treats NaNs as equal.
        template<>
        struct MultiMapAccumulatorTypeTraits<float> {
            using AccumulatorType = MultiMapAccumulator<
                float,
                util::floating_point::NaNAwareHash<float>,
                util::floating_point::NaNAwareEquals<float> >;
        };

        template<>
        struct MultiMapAccumulatorTypeTraits<double> {
            using AccumulatorType = MultiMapAccumulator<
                double,
                util::floating_point::NaNAwareHash<double>,
                util::floating_point::NaNAwareEquals<double> >;
        };

        template<>
        struct MultiMapAccumulatorTypeTraits<ComplexType> {
            using AccumulatorType = ComplexTypeMultiMapAccumulator;
        };

        template<
            typename K,
            typename AccumulatorType =
            typename MultiMapAccumulatorTypeTraits<K>::AccumulatorType>
        class MultiMapAggAggregate : public exec::Aggregate {
        public:
            explicit MultiMapAggAggregate(TypePtr resultType)
                : exec::Aggregate(std::move(resultType)) {
            }

            bool isFixedSize() const override {
                return false;
            }

            int32_t accumulatorFixedWidthSize() const override {
                return sizeof(AccumulatorType);
            }

            void addRawInput(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool mayPushdown) override {
                decodedKeys_.decode(*args[0], rows);
                decodedValues_.decode(*args[1], rows);

                rows.applyToSelected([&](vector_size_t row) {
                    // Skip null keys.
                    if (!decodedKeys_.is_null_at(row)) {
                        auto group = groups[row];
                        clearNull(group);

                        auto *accumulator = value<AccumulatorType>(group);
                        auto tracker = trackRowSize(group);
                        accumulator->insert(decodedKeys_, decodedValues_, row, *allocator_);
                    }
                });
            }

            void addIntermediateResults(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool mayPushdown) override {
                decodedMaps_.decode(*args[0], rows);
                auto map_vector = decodedMaps_.base()->template as<MapVector>();

                decodedKeys_.decode(*map_vector->mapKeys());
                decodedValueArrays_.decode(*map_vector->mapValues());

                auto *valueArrays_ = decodedValueArrays_.base()->template as<ArrayVector>();

                decodedValues_.decode(*valueArrays_->elements());

                POLLUX_CHECK_NOT_NULL(map_vector);
                rows.applyToSelected([&](vector_size_t row) {
                    if (!decodedMaps_.is_null_at(row)) {
                        auto *group = groups[row];
                        auto accumulator = value<AccumulatorType>(group);
                        clearNull(group);
                        auto decodedRow = decodedMaps_.index(row);
                        auto offset = map_vector->offsetAt(decodedRow);
                        auto size = map_vector->sizeAt(decodedRow);
                        checkNullKeys(decodedKeys_, offset, size);
                        for (auto i = offset; i < offset + size; ++i) {
                            auto numValues = valueArrays_->sizeAt(decodedValueArrays_.index(i));
                            auto valueOffset =
                                    valueArrays_->offsetAt(decodedValueArrays_.index(i));
                            accumulator->insertMultiple(
                                decodedKeys_,
                                i,
                                decodedValues_,
                                valueOffset,
                                numValues,
                                *allocator_);
                        }
                    }
                });
            }

            void addSingleGroupRawInput(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool mayPushdown) override {
                decodedKeys_.decode(*args[0], rows);
                decodedValues_.decode(*args[1], rows);

                auto *accumulator = value<AccumulatorType>(group);
                auto tracker = trackRowSize(group);

                rows.applyToSelected([&](vector_size_t row) {
                    // Skip null keys.
                    if (!decodedKeys_.is_null_at(row)) {
                        clearNull(group);
                        accumulator->insert(decodedKeys_, decodedValues_, row, *allocator_);
                    }
                });
            }

            void addSingleGroupIntermediateResults(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool mayPushdown) override {
                decodedMaps_.decode(*args[0], rows);
                auto map_vector = decodedMaps_.base()->template as<MapVector>();

                decodedKeys_.decode(*map_vector->mapKeys());
                decodedValueArrays_.decode(*map_vector->mapValues());

                auto *valueArrays_ = decodedValueArrays_.base()->template as<ArrayVector>();

                decodedValues_.decode(*valueArrays_->elements());

                auto accumulator = value<AccumulatorType>(group);

                POLLUX_CHECK_NOT_NULL(map_vector);
                rows.applyToSelected([&](vector_size_t row) {
                    if (!decodedMaps_.is_null_at(row)) {
                        clearNull(group);
                        auto decodedRow = decodedMaps_.index(row);
                        auto offset = map_vector->offsetAt(decodedRow);
                        auto size = map_vector->sizeAt(decodedRow);
                        checkNullKeys(decodedKeys_, offset, size);
                        for (auto i = offset; i < offset + size; ++i) {
                            auto numValues = valueArrays_->sizeAt(decodedValueArrays_.index(i));
                            auto valueOffset =
                                    valueArrays_->offsetAt(decodedValueArrays_.index(i));
                            accumulator->insertMultiple(
                                decodedKeys_,
                                i,
                                decodedValues_,
                                valueOffset,
                                numValues,
                                *allocator_);
                        }
                    }
                });
            }

            void extractValues(char **groups, int32_t numGroups, VectorPtr *result)
            override {
                auto map_vector = (*result)->as<MapVector>();
                POLLUX_CHECK(map_vector);
                map_vector->resize(numGroups);
                auto mapKeys = map_vector->mapKeys();
                auto mapValueArrays = map_vector->mapValues()->as<ArrayVector>();
                auto numKeys = countKeys(groups, numGroups);
                mapKeys->resize(numKeys);
                mapValueArrays->resize(numKeys);

                auto mapValues = mapValueArrays->elements();
                mapValues->resize(countValues(groups, numGroups));

                auto *rawNulls = getRawNulls(map_vector);
                vector_size_t keyOffset = 0;
                vector_size_t valueOffset = 0;

                for (int32_t i = 0; i < numGroups; ++i) {
                    char *group = groups[i];

                    if (isNull(group)) {
                        map_vector->set_null(i, true);
                        map_vector->setOffsetAndSize(i, 0, 0);
                        continue;
                    }

                    clearNull(rawNulls, i);

                    auto accumulator = value<AccumulatorType>(group);
                    auto mapSize = accumulator->size();
                    if (mapSize) {
                        map_vector->setOffsetAndSize(i, keyOffset, mapSize);
                        accumulator->extract(mapKeys, *mapValueArrays, keyOffset, valueOffset);
                    } else {
                        map_vector->setOffsetAndSize(i, 0, 0);
                    }
                }
            }

            void extractAccumulators(char **groups, int32_t numGroups, VectorPtr *result)
            override {
                extractValues(groups, numGroups, result);
            }

        protected:
            void initializeNewGroupsInternal(
                char **groups,
                melon::Range<const vector_size_t *> indices) override {
                const auto &type = resultType()->childAt(0);
                for (auto index: indices) {
                    new(groups[index] + offset_) AccumulatorType(type, allocator_);
                }
                setAllNulls(groups, indices);
            }

            void destroyInternal(melon::Range<char **> groups) override {
                for (auto *group: groups) {
                    if (isInitialized(group)) {
                        auto accumulator = value<AccumulatorType>(group);
                        accumulator->free(*allocator_);
                        destroyAccumulator<AccumulatorType>(group);
                    }
                }
            }

        private:
            static void checkNullKeys(
                const DecodedVector &keys,
                vector_size_t offset,
                vector_size_t size) {
                static const char *kNullKey = "map key cannot be null";
                if (keys.may_have_nulls()) {
                    for (auto i = offset; i < offset + size; ++i) {
                        POLLUX_USER_CHECK(!keys.is_null_at(i), kNullKey);
                    }
                }
            }

            vector_size_t countKeys(char **groups, int32_t numGroups) const {
                vector_size_t size = 0;
                for (int32_t i = 0; i < numGroups; ++i) {
                    size += value<AccumulatorType>(groups[i])->size();
                }
                return size;
            }

            vector_size_t countValues(char **groups, int32_t numGroups) const {
                vector_size_t size = 0;
                for (int32_t i = 0; i < numGroups; ++i) {
                    size += value<AccumulatorType>(groups[i])->numValues();
                }
                return size;
            }

            DecodedVector decodedKeys_;
            DecodedVector decodedValues_;
            DecodedVector decodedMaps_;
            DecodedVector decodedValueArrays_;
        };

        template<TypeKind Kind>
        std::unique_ptr<exec::Aggregate> createMultiMapAggAggregateWithCustomCompare(
            const TypePtr &resultType) {
            return std::make_unique<MultiMapAggAggregate<
                typename TypeTraits<Kind>::NativeType,
                CustomComparisonMultiMapAccumulator<Kind> > >(resultType);
        }
    } // namespace

    void registerMultiMapAggAggregate(
        const std::string &prefix,
        bool withCompanionFunctions,
        bool overwrite) {
        std::vector<std::shared_ptr<exec::AggregateFunctionSignature> > signatures{
            exec::AggregateFunctionSignatureBuilder()
            .typeVariable("K")
            .typeVariable("V")
            .returnType("map(K,array(V))")
            .intermediateType("map(K,array(V))")
            .argumentType("K")
            .argumentType("V")
            .build()
        };

        auto name = prefix + kMultiMapAgg;
        exec::registerAggregateFunction(
            name,
            std::move(signatures),
            [name](
        core::AggregationNode::Step step,
        const std::vector<TypePtr> &argTypes,
        const TypePtr &resultType,
        const core::QueryConfig & /*config*/)
        -> std::unique_ptr<exec::Aggregate> {
                const auto keyType = resultType->childAt(0);
                const auto typeKind = keyType->kind();

                if (keyType->providesCustomComparison()) {
                    return POLLUX_DYNAMIC_SCALAR_TYPE_DISPATCH(
                                createMultiMapAggAggregateWithCustomCompare,
                                typeKind,
                                resultType);
                }

                switch (typeKind) {
                    case TypeKind::BOOLEAN:
                        return std::make_unique<MultiMapAggAggregate<bool> >(resultType);
                    case TypeKind::TINYINT:
                        return std::make_unique<MultiMapAggAggregate<int8_t> >(resultType);
                    case TypeKind::SMALLINT:
                        return std::make_unique<MultiMapAggAggregate<int16_t> >(resultType);
                    case TypeKind::INTEGER:
                        return std::make_unique<MultiMapAggAggregate<int32_t> >(resultType);
                    case TypeKind::BIGINT:
                        return std::make_unique<MultiMapAggAggregate<int64_t> >(resultType);
                    case TypeKind::REAL:
                        return std::make_unique<MultiMapAggAggregate<float> >(resultType);
                    case TypeKind::DOUBLE:
                        return std::make_unique<MultiMapAggAggregate<double> >(resultType);
                    case TypeKind::TIMESTAMP:
                        return std::make_unique<MultiMapAggAggregate<Timestamp> >(
                            resultType);
                    case TypeKind::VARBINARY:
                        [[fallthrough]];
                    case TypeKind::VARCHAR:
                        return std::make_unique<MultiMapAggAggregate<StringView> >(
                            resultType);
                    case TypeKind::ARRAY:
                        [[fallthrough]];
                    case TypeKind::MAP:
                        [[fallthrough]];
                    case TypeKind::ROW:
                        return std::make_unique<MultiMapAggAggregate<ComplexType> >(
                            resultType);
                    case TypeKind::UNKNOWN:
                        return std::make_unique<MultiMapAggAggregate<int32_t> >(resultType);
                    default:
                        POLLUX_UNREACHABLE(
                            "Unexpected type {}", mapTypeKindToName(typeKind));
                }
            },
            withCompanionFunctions,
            overwrite);
    }
} // namespace kumo::pollux::aggregate::prestosql
